require 'nn'
require 'cudnn'
require 'inn'

local function createModel(opt)
	 model = torch.load('./models/googlenet-inception4e.t7')

    local CAM_conv = cudnn.SpatialConvolution(832, 1024, 3, 3, 1, 1, 1, 1, 2)
    CAM_conv.name = 'CAM_conv'
    model:add(CAM_conv)

    model:add(cudnn.ReLU(true))

    local gbvs = nn.GBVS()
    model:add(gbvs)

    model:add(nn.SpatialSumOverMap())

    model:add(nn.View(-1, 1024))
    model:add(nn.Dropout(0.5))
    
    local classifier = nn.Linear(1024, 1000)
    classifier.name = 'CAM_fc'
    model:add(classifier)

    -- print(model)
    return model
end

return createModel

-- return function(opt)

--     model = torch.load('/root/Workspace/fb.resnet.torch/models/googlenet-inception4e.t7')

--     local CAM_conv = cudnn.SpatialConvolution(832, 1024, 3, 3, 1, 1, 1, 1, 2)
--     CAM_conv.name = 'CAM_conv'
--     model:add(CAM_conv)

--     model:add(cudnn.ReLU(true))

--     local gbvs = nn.GBVS()
--     model:add(gbvs)

--     model:add(nn.SpatialSumOverMap())
--     -- model:add(cudnn.SpatialAveragePooling(14, 14, 1, 1))

--     model:add(nn.View(-1, 1024))
--     model:add(nn.Dropout(0.5))
    
--     local classifier = nn.Linear(1024, 1000)
--     classifier.name = 'CAM_fc'
--     model:add(classifier)

--     return model
-- end